package com.langll.ai.vector.service;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONObject;
import com.google.gson.Gson;
import com.google.gson.JsonObject;
import com.langll.ai.vector.config.MilvusArchive;
import io.milvus.v2.client.MilvusClientV2;
import io.milvus.v2.common.DataType;
import io.milvus.v2.common.IndexParam;
import io.milvus.v2.service.collection.request.*;
import io.milvus.v2.service.index.request.CreateIndexReq;
import io.milvus.v2.service.vector.request.DeleteReq;
import io.milvus.v2.service.vector.request.InsertReq;
import io.milvus.v2.service.vector.request.QueryReq;
import io.milvus.v2.service.vector.request.SearchReq;
import io.milvus.v2.service.vector.request.data.FloatVec;
import io.milvus.v2.service.vector.response.InsertResp;
import io.milvus.v2.service.vector.response.QueryResp;
import io.milvus.v2.service.vector.response.SearchResp;
import jakarta.annotation.PostConstruct;
import lombok.RequiredArgsConstructor;
import org.apache.commons.lang3.StringUtils;
import org.springframework.ai.document.Document;
import org.springframework.ai.document.DocumentMetadata;
import org.springframework.ai.ollama.OllamaEmbeddingModel;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.stereotype.Service;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;

/**
 * @author langll
 * @des MilvusServiceImpl
 * @date 2025/2/25 下午3:10
 */
@Service
@RequiredArgsConstructor
public class MilvusServiceImpl implements MilvusService {

    private final MilvusClientV2 milvusClientV2;
    private final OllamaEmbeddingModel ollamaEmbeddingModel;

    /**
     * 检查集合是否存在
     */
    @Override
    public void hasCollection() {
        //milvusClientV2.dropCollection(DropCollectionReq.builder().collectionName(MilvusArchive.COLLECTION_NAME).build());
        Boolean b = milvusClientV2.hasCollection(HasCollectionReq.builder().collectionName(MilvusArchive.COLLECTION_NAME).build());
        if (!b) {
            this.createCollection();
        }
    }

    /**
     * 创建集合
     */
    public void createCollection() {
        // 创建字段
        CreateCollectionReq.CollectionSchema schema = milvusClientV2.createSchema()
                // 创建主键字段
                .addField(AddFieldReq.builder()
                        // 字段名
                        .fieldName(MilvusArchive.Field.ID)
                        // 字段描述
                        .description("主键ID")
                        // 字段类型
                        .dataType(DataType.Int64)
                        // 是否为主键
                        .isPrimaryKey(true)
                        // 设置主键自增
                        .autoID(true)
                        .build())
                .addField(AddFieldReq.builder()
                        // 字段名
                        .fieldName(MilvusArchive.Field.FILE_NAME)
                        // 字段描述
                        .description("文件名")
                        // 字段类型
                        .dataType(DataType.VarChar)
                        // 设置字段为可空
                        .isNullable(true)
                        .build())
                // 创建特征向量字段
                .addField(AddFieldReq.builder()
                        // 字段名
                        .fieldName(MilvusArchive.Field.FEATURE)
                        // 字段描述
                        .description("特征向量")
                        // 字段类型
                        .dataType(DataType.FloatVector)
                        // 设置向量维度
                        .dimension(MilvusArchive.FEATURE_DIM)
                        .build())
                .addField(AddFieldReq.builder()
                        // 字段名
                        .fieldName(MilvusArchive.Field.TEXT)
                        // 字段描述
                        .description("文本")
                        // 字段类型
                        .dataType(DataType.VarChar)
                        // 设置字段为可空
                        .isNullable(true)
                        .build())
                .addField(AddFieldReq.builder()
                        // 字段名
                        .fieldName(MilvusArchive.Field.METADATA)
                        // 字段描述
                        .description("元数据")
                        // 字段类型
                        .dataType(DataType.VarChar)
                        // 设置字段为可空
                        .isNullable(true)
                        .build());
        // 创建集合
        CreateCollectionReq collectionReq = CreateCollectionReq.builder()
                // 集合名称
                .collectionName(MilvusArchive.COLLECTION_NAME)
                // 集合描述
                .description("自定义知识库")
                // 集合字段
                .collectionSchema(schema)
                // 分片数量
                .numShards(MilvusArchive.SHARDS_NUM)
                .build();
        milvusClientV2.createCollection(collectionReq);

        // 创建索引
        IndexParam indexParam = IndexParam.builder()
                // 索引字段名
                .fieldName(MilvusArchive.Field.FEATURE)
                // 索引类型
                .indexType(IndexParam.IndexType.IVF_FLAT)
                // 索引距离度量
                .metricType(IndexParam.MetricType.COSINE)
                .build();
        CreateIndexReq createIndexReq = CreateIndexReq.builder()
                .collectionName(MilvusArchive.COLLECTION_NAME)
                .indexParams(Collections.singletonList(indexParam))
                .build();

        milvusClientV2.createIndex(createIndexReq);
    }

    /**
     * 插入数据
     */
    @Override
    public InsertResp insert(float[] vectorParam, String text, String metadata, String fileName) {
        // 校验集合是否存在
        this.hasCollection();
        JsonObject jsonObject = new JsonObject();
        // 数组转换成JsonElement
        jsonObject.add(MilvusArchive.Field.FEATURE, new Gson().toJsonTree(vectorParam));
        jsonObject.add(MilvusArchive.Field.TEXT, new Gson().toJsonTree(text));
        jsonObject.add(MilvusArchive.Field.METADATA, new Gson().toJsonTree(metadata));
        jsonObject.add(MilvusArchive.Field.FILE_NAME, new Gson().toJsonTree(fileName));
        InsertReq insertReq = InsertReq.builder()
                // 集合名称
                .collectionName(MilvusArchive.COLLECTION_NAME)
                .data(Collections.singletonList(jsonObject))
                .build();

        return milvusClientV2.insert(insertReq);
    }

    /**
     * 批量插入数据
     *
     * @param vectorParam 向量参数
     * @param text        文本
     * @param metadata    元数据
     * @param fileName    文件名
     */
    @Override
    public InsertResp batchInsert(List<float[]> vectorParam, List<String> text, List<String> metadata, List<String> fileName) {
        if (vectorParam.size() == text.size() && vectorParam.size() == metadata.size() && vectorParam.size() == fileName.size()) {
            List<JsonObject> jsonObjects = new ArrayList<>();
            for (int i = 0; i < vectorParam.size(); i++) {
                JsonObject jsonObject = new JsonObject();
                // 数组转换成JsonElement
                jsonObject.add(MilvusArchive.Field.FEATURE, new Gson().toJsonTree(vectorParam.get(i)));
                jsonObject.add(MilvusArchive.Field.TEXT, new Gson().toJsonTree(text.get(i)));
                jsonObject.add(MilvusArchive.Field.METADATA, new Gson().toJsonTree(metadata.get(i)));
                jsonObject.add(MilvusArchive.Field.FILE_NAME, new Gson().toJsonTree(fileName.get(i)));
                jsonObjects.add(jsonObject);
            }
            InsertReq insertReq = InsertReq.builder()
                    // 集合名称
                    .collectionName(MilvusArchive.COLLECTION_NAME)
                    .data(jsonObjects)
                    .build();
            return milvusClientV2.insert(insertReq);
        }
        return null;
    }

    /**
     * 搜索数据
     *
     * @param vectorParam 向量参数
     */
    @Override
    public SearchResp search(float[] vectorParam) {
        this.loadCollection();
        FloatVec floatVec = new FloatVec(vectorParam);
        SearchReq searchReq = SearchReq.builder()
                // 集合名称
                .collectionName(MilvusArchive.COLLECTION_NAME)
                // 搜索距离度量
                .metricType(IndexParam.MetricType.COSINE)
                // 搜索向量
                .data(Collections.singletonList(floatVec))
                // 搜索字段
                .annsField(MilvusArchive.Field.FEATURE)
                // 返回字段
                .outputFields(Arrays.asList(MilvusArchive.Field.ID, MilvusArchive.Field.TEXT, MilvusArchive.Field.METADATA, MilvusArchive.Field.FILE_NAME))
                // 搜索数量
                .topK(5)
                .build();
        return milvusClientV2.search(searchReq);
    }

    @PostConstruct
    public void loadCollection() {
        this.hasCollection();
        // 查询集合状态
        if (!milvusClientV2.getLoadState(GetLoadStateReq.builder().collectionName(MilvusArchive.COLLECTION_NAME).build())) {
            milvusClientV2.loadCollection(LoadCollectionReq.builder().collectionName(MilvusArchive.COLLECTION_NAME).build());
        }
    }

    public void releaseCollection() {
        milvusClientV2.releaseCollection(ReleaseCollectionReq.builder().collectionName(MilvusArchive.COLLECTION_NAME).build());
    }

    /**
     * 删除数据
     *
     * @param ids id
     */
    @Override
    public void delete(String... ids) {
        milvusClientV2.delete(DeleteReq.builder().collectionName(MilvusArchive.COLLECTION_NAME).ids(Arrays.asList(ids)).build());
    }

    @Override
    public void delete(String filter) {
        milvusClientV2.delete(DeleteReq.builder().collectionName(MilvusArchive.COLLECTION_NAME).filter(filter).build());
    }

    /**
     * 搜索数据
     *
     * @param request 搜索请求
     */
    @Override
    public List<Document> search(SearchRequest request) {
        String query = request.getQuery();
        float[] embed = ollamaEmbeddingModel.embed(query);
        List<List<SearchResp.SearchResult>> searchResults = this.search(embed).getSearchResults();
        if (!searchResults.isEmpty()) {
            return searchResults.get(0).stream().filter(searchResult -> searchResult.getScore() >= request.getSimilarityThreshold()).map(searchResult -> {
                String docId = searchResult.getId().toString();
                String content = searchResult.getEntity().get(MilvusArchive.Field.TEXT).toString();
                JSONObject metadata = JSON.parseObject(searchResult.getEntity().get(MilvusArchive.Field.METADATA).toString());
                metadata.put(DocumentMetadata.DISTANCE.value(), 1.0F - searchResult.getScore());
                return Document.builder().id(docId).text(content).metadata(metadata.getInnerMap()).score((double) searchResult.getScore()).build();
            }).toList();
        }
        return List.of();
    }

    /**
     * 查询数据
     *
     * @param page     页码
     * @param size     每页数量
     * @param fileName 文件名模糊搜索
     */
    @Override
    public List<QueryResp.QueryResult> query(Integer page, Integer size, String fileName) {
        // 查询集合状态并加载到内存
        this.loadCollection();
        QueryReq queryReq = QueryReq.builder()
                .collectionName(MilvusArchive.COLLECTION_NAME)
                .offset((long) (page - 1) * size)
                .limit(size)
                .outputFields(List.of(MilvusArchive.Field.FILE_NAME, MilvusArchive.Field.TEXT))
                .build();
        if (StringUtils.isNotBlank(fileName)) {
            queryReq.setFilter((String.format("%s LIKE \"%%%s%%\"", MilvusArchive.Field.FILE_NAME, fileName)));
        }
        QueryResp query = milvusClientV2.query(queryReq);
        return query.getQueryResults();
    }
}
